{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# 3.4 A Nonlinear conservation law: shallow water equation in 2D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "We consider the shallow water equations as an example of a nonlinear conservation law, i.e. we consider\n",
    "\n",
    "$$\n",
    "  \\partial_t \\mathbf{U} + \\operatorname{div}(\\mathbf{F} (\\mathbf{U} )) = 0 \\qquad in \\qquad \\Omega \\times[0,T],\n",
    "$$\n",
    "\n",
    "with \n",
    "\n",
    "$$\n",
    "\\mathbf{U} = (h, \\mathbf{w}) = (h, hu, hv) = (\\mathbf{u}_1, \\mathbf{u}_2, \\mathbf{u}_3)\n",
    "$$\n",
    "\n",
    "and \n",
    "\n",
    "$$\n",
    "  \\mathbf{F}(\\mathbf{U}) = \\left( \\begin{array}{c@{}c@{\\qquad}c@{}c} h u & & h v & \\\\ h u^2 &+ \\frac12 g h^2 & h u v & \\\\ h u v &  & h v^2 & + \\frac12 g h^2 \\end{array} \\right) \n",
    "  = \\left( \\begin{array}{cc} \\mathbf{u}_2 & \\mathbf{u}_3 \\\\ \\frac{\\mathbf{u}_2^2}{\\mathbf{u}_1} + \\frac12 g \\mathbf{u}_1^2 & \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1} \\\\ \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1} & \\frac{\\mathbf{u}_3^2}{\\mathbf{u}_1} + \\frac12 g \\mathbf{u}_1^2 \\end{array} \\right)  \n",
    "  = \\left( \\begin{array}{cc} h \\mathbf{w}^T \\\\ h \\mathbf{w} \\otimes \\mathbf{w} + \\frac12 g h^2 \\mathbf{I} \\end{array} \\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "## Jacobian of the flux for shallow water:\n",
    "\n",
    "\\begin{align*}\n",
    "\\mathbf{A}_1 & = \\left(\\begin{array}{ccc} 0 & 1 & 0 \\\\ - \\frac{\\mathbf{u}_2^2}{\\mathbf{u}_1^2} + g \\mathbf{u}_1 & 2 \\frac{\\mathbf{u}_2}{\\mathbf{u}_1} & 0 \\\\ - \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1^2} & \\frac{\\mathbf{u}_3}{\\mathbf{u}_1} & \\frac{\\mathbf{u}_2}{\\mathbf{u}_1} \\end{array} \\right) = \\left( \\begin{array}{ccc} 0 & 1 & 0 \\\\ - u^2 + g h & 2 u & 0 \\\\ - u v & v & u \\end{array} \\right) \n",
    "\\end{align*}\n",
    "\n",
    "\\begin{align*}\n",
    "\\mathbf{A}_2 & = \\left( \\begin{array}{ccc} 0 & 0 & 1 \\\\ - \\frac{\\mathbf{u}_2\\mathbf{u}_3}{\\mathbf{u}_1^2} & \\frac{\\mathbf{u}_3}{\\mathbf{u}_1} & \\frac{\\mathbf{u}_2}{\\mathbf{u}_1} \\\\ - \\frac{\\mathbf{u}_3^2}{\\mathbf{u}_1^2} + g \\mathbf{u}_1  & 0 & 2\\frac{\\mathbf{u}_3}{\\mathbf{u}_1} \\end{array} \\right) = \\left( \\begin{array}{ccc} 0 & 0 & 1 \\\\ - uv & v & u \\\\ - v^2 + gh & 0 & 2 v \\end{array} \\right)\n",
    "\\end{align*}\n",
    "\n",
    "\\begin{align*}\n",
    "\\mathbf{A}(\\mathbf{u}, \\mathbf{n}) & = n_1 \\mathbf{A}_1 + n_2 \\mathbf{A}_2 = \\left( \\begin{array}{ccc} 0 & n_1 & n_2 \\\\\n",
    "- u \\alpha - g h n_1  & \\alpha + u n_1 & u n_2 \\\\\n",
    "- v \\alpha - g h n_2  & v n_1 & \\alpha + v n_2 \\end{array} \\right), \\quad \\text{ with } \\alpha = (\\mathbf{u}, \\mathbf{n}),\n",
    "\\end{align*}\n",
    "\n",
    "spectrum:\n",
    "\n",
    "$$\n",
    "\\rho( \\mathbf{A}(\\mathbf{u}, \\mathbf{n}) ) = \\{ \\alpha + \\sqrt{g h}, \\alpha, \\alpha - \\sqrt{g h} \\} \n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "from ngsolve import * # sloppy\n",
    "import netgen.gui"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### The dam break problem (geometry)\n",
    "\n",
    "![alt](geomsketch.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "from netgen.geom2d import SplineGeometry\n",
    "geo = SplineGeometry()\n",
    "\n",
    "pnts =[ (-12,-5), (-7,-5), (-5,-5), (-3,-5), (-3,-3), \n",
    "        (-3,-1), (-1,-1), ( 0,-1), ( 1,-1), ( 3,-1), \n",
    "        ( 3,-3), ( 3,-5), ( 5,-5), ( 7,-5), ( 12,-5), \n",
    "        ( 12, 5), ( 7, 5), ( 5, 5), ( 3, 5), ( 3, 3), \n",
    "        ( 3, 1), ( 1, 1), ( 0, 1), (-1, 1), (-3, 1), \n",
    "        (-3, 3), (-3, 5), (-5, 5), (-7, 5), (-12, 5)]\n",
    "pnts = [geo.AppendPoint(*pnt) for pnt in pnts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "curves = [[[\"line\",0,1],\"wall\",1, 0], [[\"line\",1,2],\"wall\",1, 0],\n",
    "          [[\"spline3\",2,3,4],\"wall\",1, 0],[[\"spline3\",4,5,6],\"wall\",1, 0],[[\"line\",6,7],\"wall\",1, 0],\n",
    "          [[\"line\",7,22],\"dam\",1, 2],     # <--- dam interface\n",
    "          [[\"line\",7,8],\"wall\",2, 0],[[\"spline3\",8,9,10],\"wall\",2, 0],\n",
    "          [[\"spline3\",10,11,12],\"wall\",2, 0],[[\"line\",12,13],\"wall\",2, 0],[[\"line\",13,14],\"wall\",2, 0],\n",
    "          [[\"line\",14,15],\"wall\",2, 0], # <--- right boundary\n",
    "          [[\"line\",15,16],\"wall\",2, 0],[[\"line\",16,17],\"wall\",2, 0],\n",
    "          [[\"spline3\",17,18,19],\"wall\",2, 0],[[\"spline3\",19,20,21],\"wall\",2, 0],\n",
    "          [[\"line\",21,22],\"wall\",2, 0],[[\"line\",22,23],\"wall\",1, 0],\n",
    "          [[\"spline3\",23,24,25],\"wall\",1, 0],[[\"spline3\",25,26,27],\"wall\",1, 0],\n",
    "          [[\"line\",27,28],\"wall\",1, 0],[[\"line\",28,29],\"wall\",1, 0],\n",
    "          [[\"line\",29,0],\"wall\",1, 0]]   # <--- left boundary\n",
    "\n",
    "for c,bc,l,r in curves:\n",
    "    geo.Append(c,bc=bc,leftdomain=l, rightdomain=r) \n",
    "geo.SetMaterial(1,\"upperlevel\")\n",
    "geo.SetMaterial(2,\"lowerlevel\")    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "mesh = Mesh(geo.GenerateMesh(maxh=2))\n",
    "mesh.Curve(3)\n",
    "Draw(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Vectorial (`dim=3`) approximation space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "order = 2\n",
    "fes = L2(mesh,order=order,dim=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### initial and boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "U,V = fes.TnT() # \"Trial\" and \"Test\" function\n",
    "h, hu, hv = U\n",
    "\n",
    "# initial conditions\n",
    "h0mat = {\"upperlevel\" : 10, \"lowerlevel\" : 2}\n",
    "U0 = CoefficientFunction((CoefficientFunction([h0mat[mat] for mat in mesh.GetMaterials()]),0,0))\n",
    "\n",
    "# boundary conditions\n",
    "hbndreg = CoefficientFunction([{\"wall\" : h, \"dam\" : 0}[rg] for rg in mesh.GetBoundaries()])\n",
    "hubndreg = CoefficientFunction([{\"wall\" : -hu, \"dam\" : 0}[rg] for rg in mesh.GetBoundaries()])\n",
    "hvbndreg = CoefficientFunction([{\"wall\" : -hv, \"dam\" : 0}[rg] for rg in mesh.GetBoundaries()])\n",
    "\n",
    "Ubnd = CoefficientFunction((hbndreg,hubndreg,hvbndreg))\n",
    "\n",
    "# constant for gravitational force\n",
    "g=1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Flux definition and numerical flux choice (Lax-Friedrich)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "$$\n",
    "  \\mathbf{F}(\\mathbf{U}) = \\left( \\begin{array}{c@{}c@{\\qquad}c@{}c} h u & & h v & \\\\ h u^2 &+ \\frac12 g h^2 & h u v & \\\\ h u v &  & h v^2 & + \\frac12 g h^2 \\end{array} \\right) \n",
    "  = \\left( \\begin{array}{cc} \\mathbf{u}_2 & \\mathbf{u}_3 \\\\ \\frac{\\mathbf{u}_2^2}{\\mathbf{u}_1} + \\frac12 g \\mathbf{u}_1^2 & \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1} \\\\ \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1} & \\frac{\\mathbf{u}_3^2}{\\mathbf{u}_1} + \\frac12 g \\mathbf{u}_1^2 \\end{array} \\right)  \n",
    "  = \\left( \\begin{array}{cc} h \\mathbf{w}^T \\\\ h \\mathbf{w} \\otimes \\mathbf{w} + \\frac12 g h^2 \\mathbf{I} \\end{array} \\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "def F(U):\n",
    "    h, hvx, hvy = U\n",
    "    vx = hvx/h\n",
    "    vy = hvy/h\n",
    "    return CoefficientFunction(((hvx,hvy),\n",
    "                                (hvx*vx + 0.5*g*h**2, hvx*vy),\n",
    "                                (hvx*vy, hvy*vy + 0.5*g*h**2)),dims=(3,2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "$$\n",
    "  \\hat{\\mathbf{F}}(\\mathbf{U}) \\cdot \\mathbf{n} = \\frac{(\\mathbf{F}(\\mathbf{U}^l)+\\mathbf{F}(\\mathbf{U}^r)}{2} \\cdot \\mathbf{n} + \\frac{|\\lambda|}{2} (\\mathbf{U}^l - \\mathbf{U}^r), \\quad (\\cdot^l: \\text{current element}, \\quad \\cdot^r: \\text{neighbor element})\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "n = specialcf.normal(mesh.dim)\n",
    "\n",
    "def Max(u,v):\n",
    "    return IfPos(u-v,u,v)\n",
    "def Fmax(A,B): # max. characteristic speed:\n",
    "    ha, hua, hva = A\n",
    "    hb, hub, hvb = B\n",
    "    vnorma = sqrt(hua**2+hva**2)/ha\n",
    "    vnormb = sqrt(hub**2+hvb**2)/hb\n",
    "    return Max(vnorma+sqrt(g*A[0]),vnormb+sqrt(g*B[0]))\n",
    "\n",
    "def Fhatn(U): # numerical flux\n",
    "    Uhat = U.Other(bnd=Ubnd)\n",
    "    return (0.5*F(U)+0.5*F(Uhat))*n + Fmax(U,Uhat)/2*(U-Uhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### DG formulation\n",
    "\n",
    "We recall that a `BilinearForm` is allowed to be nonlinear in the first argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "def DGBilinearForm(fes,F,Fhatn,Ubnd):\n",
    "    a = BilinearForm(fes, nonassemble=True)    \n",
    "    a += - InnerProduct(F(U),Grad(V)) * dx \n",
    "    a += InnerProduct(Fhatn(U),V) * dx(element_boundary=True)\n",
    "    return a\n",
    "\n",
    "a = DGBilinearForm(fes,F,Fhatn,Ubnd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "#### Simple fix to deal with shocks: artificial diffusion:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from DGdiffusion import AddArtificialDiffusion\n",
    "\n",
    "artvisc = Parameter(1.0)\n",
    "if order > 0:\n",
    "    AddArtificialDiffusion(a,Ubnd,artvisc,compile=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Visualization of solution quantities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "gfu = GridFunction(fes)\n",
    "gfh, gfhu, gfhv = gfu\n",
    "gfvu = gfhu/gfh\n",
    "gfvv = gfhv/gfh\n",
    "momentum = CoefficientFunction((gfhu,gfhv))\n",
    "velocity = CoefficientFunction((gfvu,gfvv))\n",
    "gfu.Set(U0)\n",
    "Draw(momentum,mesh,\"mom\")\n",
    "Draw(velocity,mesh,\"vel\")\n",
    "Draw(gfh,mesh,\"h\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Explicit Euler time stepping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "def TimeLoop(a,gfu,dt,T,nsamplings=100):\n",
    "    #gfu.Set(U0)\n",
    "    res = gfu.vec.CreateVector()\n",
    "    fes = a.space\n",
    "    t = 0; i = 0\n",
    "    nsteps = int(ceil(T/dt))    \n",
    "    invma = fes.Mass(1).Inverse() @ a.mat\n",
    "    with TaskManager():\n",
    "        while t <= T - 0.5*dt:\n",
    "            res = invma * gfu.vec\n",
    "            gfu.vec.data -= dt * res\n",
    "            t += dt\n",
    "            if (i+1) % int(nsteps/nsamplings) == 0:\n",
    "                Redraw()\n",
    "            i+=1\n",
    "            print(\"\\rt = {:.10}\".format(t),end=\"\")\n",
    "    Redraw()     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "TimeLoop(a,gfu,dt=0.0004,T=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "You may play around with this example. \n",
    "\n",
    "* change the artificial diffusion parameter: How does it influence the solution\n",
    "* change boundary conditions: left boundary -> fixed height and non-reflecting\n",
    "* change the initial heights\n",
    "* introduce a (circular) obstacle below the dam\n",
    "* Generate output to create a video: To this end take a look at the final unit of this section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": true,
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "%%HTML\n",
    "<video width=\"600\" height=\"400\" controls>\n",
    "    <source src=\"../../_static/shallow2D.mov\">\n",
    "</video>"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
